package com.ai.service.impl;

import cn.hutool.core.io.FileUtil;
import cn.hutool.core.util.RuntimeUtil;
import cn.hutool.core.util.StrUtil;
import com.ai.service.VectorStoreService;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.reader.TextReader;
import org.springframework.ai.reader.tika.TikaDocumentReader;
import org.springframework.ai.transformer.splitter.TextSplitter;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.SimpleVectorStore;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.filter.Filter;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Bean;
import org.springframework.core.io.FileSystemResource;
import org.springframework.core.io.InputStreamResource;
import org.springframework.core.io.Resource;
import org.springframework.stereotype.Service;
import org.springframework.web.multipart.MultipartFile;

import java.io.File;
import java.io.IOException;
import java.time.Duration;
import java.util.ArrayList;
import java.util.List;
import java.util.Timer;
import java.util.TimerTask;




@Service
@Slf4j
public class VectorStoreServiceImpl implements VectorStoreService {

    private static final String VECTOR_STORE_DOCUMENT_ID = "documentId";

    @Autowired
    private EmbeddingModel embeddingModel;

    @Autowired
    private VectorStore vectorStore;

    @Autowired
    private DocumentSplitterService documentSplitterService;


    /*
     *  简单保存
     */
    @Override
    public int simpleVectorStore() {
        String content = "这是一段测试文本";
        Document document = new Document(content);
        // 文件存储
        VectorStore vectorStore = buildSimpleVectorStore(embeddingModel);

        // 存储数据 业务关联
        document.getMetadata().put(VECTOR_STORE_DOCUMENT_ID, "c7e8354c-af69-721b-dac0-cc0d8833b2b4");
        vectorStore.add(List.of(document));
        log.info("vector store success...");
        return 1;
    }


    private VectorStore buildSimpleVectorStore(EmbeddingModel embeddingModel) {
        SimpleVectorStore vectorStore = SimpleVectorStore.builder(embeddingModel).build();
        // 本地存储文件 FileUtil.getUserHomePath()
        File file = new File(StrUtil.format("{}/data_store/simple_{}.json", "/data", embeddingModel.getClass().getSimpleName()));
        log.info("vector store file path: {}", file.getAbsolutePath());
        if (!file.exists()) {
            FileUtil.mkParentDirs(file);
            try {
                boolean newFile = file.createNewFile();
            } catch (IOException e) {
                log.info("file create error...");
            }
        } else if (file.length() > 0) {
            vectorStore.load(file);
        }
        // 定时器
        Timer timer = new Timer("SimpleVectorStoreTimer-" + file.getAbsolutePath());
        log.info("timer task start...");
        timer.scheduleAtFixedRate(new TimerTask() {
            @Override
            public void run() {
                vectorStore.save(file);
            }
        }, Duration.ofMinutes(1).toMillis(), Duration.ofMinutes(1).toMillis());
        // 持久化
        RuntimeUtil.addShutdownHook(() -> vectorStore.save(file));
        return vectorStore;
    }


    /**
     * 保存文档到向量数据库 自动处理
     * @param file MultipartFile
     * @return Boolean
     */
    public Boolean saveDocumentToVectorStore(MultipartFile file) throws IOException {
        // 1. 读取并解析文档
        Resource resource = new InputStreamResource(file.getInputStream());
        TikaDocumentReader tikaDocumentReader = new TikaDocumentReader(resource);
        List<Document> originalDocuments = tikaDocumentReader.get();

        // 2. 分割文本
        TextSplitter textSplitter = new TokenTextSplitter();
        List<Document> splitDocuments = textSplitter.apply(originalDocuments);

        // 3. (可选) 添加元数据，例如记录来源文件
        for (Document doc : splitDocuments) {
            doc.getMetadata().put("source_file", file.getOriginalFilename());
        }

        // 4. 存入向量数据库（通常会自动调用 EmbeddingModel 生成向量）
        vectorStore.add(splitDocuments);

        return true;
    }



    /**
     * 将上传的文件处理后存入向量数据库  自定义实现
     * @param file 上传的文件
     * @return 操作是否成功
     */
    public Boolean customerSaveDocumentToVectorStore(MultipartFile file) {
        try {

            //// 1. 加载文档
            //Resource resource = new FileSystemResource("path/to/document.txt");
            //TextReader textReader = new TextReader(resource);
            //List<Document> documents = textReader.get();
            //// 2. 分割文档为块
            //TokenTextSplitter splitter = new TokenTextSplitter();
            //List<Document> chunks = splitter.apply(documents);
            //// 3. 将块添加到向量存储
            //vectorStore.add(chunks);
            //// 现在你可以使用向量存储进行检索
            //List<Document> results = vectorStore.similaritySearch("查询文本");

            // 1. 分割文档 chunks
            List<Document> splitDocuments = documentSplitterService.splitDocumentToDocuments(file);

            // 2. 添加元数据，例如记录来源文件
            for (Document doc : splitDocuments) {
                doc.getMetadata().put("source_file", file.getOriginalFilename());
                doc.getMetadata().put("file_size", file.getSize());
                doc.getMetadata().put("upload_time", System.currentTimeMillis());
            }

            // 3. 存入向量数据库（会自动调用 EmbeddingModel 生成向量）:cite[3]
            vectorStore.add(splitDocuments);

            // 创建保存文件路径
            //File fileToSave = new File(StrUtil.format("{}/data_store/simple_{}_{}.json",
            //        "/data", embeddingModel.getClass().getSimpleName(), System.currentTimeMillis()));
            //// 确保父目录存在
            //FileUtil.mkParentDirs(fileToSave);

            //- 每次调用 customerSaveDocumentToVectorStore 时：
            //            //  使用的是 Spring 容器中单例的 vectorStore Bean
            //            //  这个 vectorStore 实例在应用启动时就已经加载了所有已存在的向量数据文件
            //            //  新上传的文档被添加到这个已有的 vectorStore 实例中
            //            //  保存时会将整个 vectorStore 的内容（包括之前的数据）保存到文件中
            //            // VectorStoreConfig中配置了vectorStore实例
            // 4. 如果vectorStore是SimpleVectorStore的实例，则可以保存到文件（否则保存内存中）
            //if (vectorStore instanceof SimpleVectorStore simpleVectorStore) {
            //    // 保存到文件(所有文件向量内容会叠加到一个文件，不建议)
            //    simpleVectorStore.save(fileToSave);
            //}

            // 每个文件单独存储 + 统一索引
            //SimpleVectorStore simpleVectorStore = SimpleVectorStore.builder(embeddingModel).build();
            ////simpleVectorStore.add(splitDocuments);
            //simpleVectorStore.save(fileToSave);

            return true;
        } catch (Exception e) {
            throw new RuntimeException("保存到向量数据库失败: " + e.getMessage(), e);
        }
    }

    /**
     * 相似性搜索
     * @param query 查询文本
     * @param topK 返回最相似的结果数量
     * @return 相似的文档列表
     */
    public List<Document> similaritySearch(String query, int topK) {
        String sourceFile = "file.pdf";

        // 检索多个文档
        List<String> sourceFiles = new ArrayList<>();
        // 添加多个文件名 ...
        sourceFiles.add("file1.pdf");
        String filterExpression = "source_file in ['" +
                String.join("','", sourceFiles) + "']";

        // 单文件
        //SearchRequest searchRequest = SearchRequest.builder()
        //        .query(query)
        //        .topK(topK).similarityThreshold(0.7) // 相似度阈值
        //        .filterExpression("source_file == '" + sourceFile + "'")
        //        .build(); // 关键：元数据过滤

        // 多文件
        SearchRequest searchRequest = SearchRequest.builder()
                .query(query)
                .topK(topK)
                .similarityThreshold(0.7) // 相似度阈值
                .filterExpression(filterExpression)
                .build(); // 关键：元数据过滤


        return vectorStore.similaritySearch(searchRequest);
    }


}
